#!/bin/bash

# 脚本名称：switch_rootfs.sh

# 功能说明：
# 本脚本用于挂载必要的文件系统并进入 chroot 环境。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 mount、chroot 等）。
# 2. 根据用户输入的根文件系统目录，挂载必要的文件系统。
# 3. 进入 chroot 环境。
# 4. 退出 chroot 环境后，卸载文件系统。

set -euo pipefail

# 检查并安装软件包
install_package() {
    local pkg="$1"
    
    if command -v apt-get &>/dev/null; then
        sudo apt-get update
        sudo apt-get install -y "$pkg"
    else
        echo "Unsupported package manager. Please install '$pkg' manually."
        exit 1
    fi
}

# 检查命令是否存在，否则尝试安装
check_command() {
    local cmd="$1"
    local pkg="${2:-$1}"  # 如果未提供包名，使用命令名作为包名

    if ! command -v "$cmd" &>/dev/null; then
        echo "Command '$cmd' not found. Installing package '$pkg'..."
        install_package "$pkg"
    fi
}

# 检查并安装所需的软件包
check_dependencies() {
    local required_commands=(
        "mount:mount"
        "chroot:coreutils"
        "qemu-aarch64-static:qemu-user-static"
        "debootstrap:debootstrap"
        "binfmt-support:binfmt-support"
    )

    for cmd_pkg in "${required_commands[@]}"; do
        IFS=':' read -r cmd pkg <<< "$cmd_pkg"
        check_command "$cmd" "$pkg"
    done
}

# 挂载必要的文件系统
mount_filesystems() {
    local root="$1"
    echo "Mounting necessary filesystems for $root..."
    sudo mount --bind /proc "$root/proc"
    sudo mount --bind /sys "$root/sys"
    sudo mount --bind /dev "$root/dev"
    sudo mount --bind /dev/pts "$root/dev/pts"
    sudo mount --bind /run "$root/run"
}

# 卸载挂载的文件系统
umount_filesystems() {
    local root="$1"
    echo "Unmounting filesystems for $root..."
    sudo umount -l "$root/proc"
    sudo umount -l "$root/sys"
    sudo umount -l "$root/dev/pts"
    sudo umount -l "$root/dev"
    sudo umount -l "$root/run"
}

# 检查并复制 QEMU 静态二进制文件
prepare_qemu() {
    local root="$1"
    local arch="$2"
    local qemu_bin="/usr/bin/qemu-${arch}-static"

    # 如果是本地架构，则不需要 QEMU
    if [ "$(uname -m)" = "$arch" ]; then
        return ""
    fi

    if [ ! -f "$qemu_bin" ]; then
        echo "Error: QEMU binary '$qemu_bin' not found. Installing qemu-user-static..."
        check_command "qemu-${arch}-static" "qemu-user-static"
    fi

    echo "Copying QEMU binary to $root..."
    sudo cp "$qemu_bin" "$root$qemu_bin"
    echo "$qemu_bin" # 返回 QEMU 文件路径
}

# 删除临时的 QEMU 静态二进制文件
cleanup_qemu() {
    local root="$1"
    local qemu_bin="$2"

    if [ -f "$root$qemu_bin" ]; then
        echo "Removing temporary QEMU binary $qemu_bin from $root..."
        sudo rm -f "$root$qemu_bin"
    fi
}

# 切换根文件系统
switch_root() {
    local root="$1"
    local arch="$2"
    local qemu_bin=""

    # 准备 QEMU（如有需要）
    qemu_bin=$(prepare_qemu "$root" "$arch")

    # 挂载必要的文件系统
    mount_filesystems "$root"

    # 使用 chroot 进入新根
    echo "Switching root to $root with architecture $arch..."
    sudo chroot "$root" /bin/bash --rcfile <(cat ~/.bashrc; echo 'PS1="(chroot)> "')

    # 退出后卸载文件系统
    echo "Exiting chroot environment..."
    umount_filesystems "$root"

    # 清理 QEMU 静态二进制文件
    if [ -n "$qemu_bin" ]; then
        cleanup_qemu "$root" "$qemu_bin"
    fi
}

# 获取支持的架构列表
get_supported_architectures() {
    local qemu_path="/usr/bin"
    local archs=()
    
    for qemu_bin in "$qemu_path"/qemu-*-static; do
        if [ -f "$qemu_bin" ]; then
            local arch=$(basename "$qemu_bin" | sed 's/qemu-\(.*\)-static/\1/')
            archs+=("$arch")
        fi
    done
    
    echo "${archs[@]}"
}

# 处理 chroot 环境并进行清理
chroot_env() {
    local rootfs_dir=$1
    
    # 确保根文件系统目录存在且有效
    if [ ! -d "$rootfs_dir" ]; then
        echo "错误：根文件系统目录 $rootfs_dir 不存在。"
        exit 1
    fi
    
    # 挂载必要的文件系统
    mount_filesystems "$rootfs_dir"

    # 确保在 chroot 会话结束后卸载文件系统
    trap "umount_filesystems $rootfs_dir" EXIT

    echo "进入 chroot 环境..."
    
    # 进入 chroot 环境
    sudo chroot "$rootfs_dir" /bin/bash
}

# 主函数
main() {
    if [ "$#" -lt 1 ]; then
        echo "Usage: $0 <root_filesystem_path> [architecture]"
        echo "Example: $0 /mnt/rootfs aarch64"
        echo "Supported architectures: $(get_supported_architectures)"
        exit 1
    fi

    local root="$1"
    local arch="${2:-}"

    # 检查依赖
    check_dependencies

    if [ ! -d "$root" ]; then
        echo "Error: Root filesystem path '$root' does not exist."
        exit 1
    fi

    if [ -n "$arch" ]; then
        # 检查架构是否支持
        local supported=0
        for supported_arch in $(get_supported_architectures); do
            if [ "$arch" = "$supported_arch" ]; then
                supported=1
                break
            fi
        done

        if [ "$supported" -eq 1 ]; then
            switch_root "$root" "$arch"
        else
            echo "Error: Unsupported architecture '$arch'."
            echo "Supported architectures: $(get_supported_architectures)"
            exit 1
        fi
    else
        chroot_env "$root"
    fi
}

main "$@"